{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tune Tutorial"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"tune.png\" alt=\"Tune Logo\" width=\"400\"/>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tune is a scalable framework for model training and hyperparameter search with a focus on deep learning and deep reinforcement learning.\n",
    "\n",
    "**Code**: https://github.com/ray-project/ray/tree/master/python/ray/tune\n",
    "\n",
    "**Examples**: https://github.com/ray-project/ray/tree/master/python/ray/tune/examples\n",
    "\n",
    "**Documentation**: http://ray.readthedocs.io/en/latest/tune.html\n",
    "\n",
    "**Mailing List** https://groups.google.com/forum/#!forum/ray-dev"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Overview\n",
    "\n",
    "Tuning hyperparameters is often the most expensive part of the machine learning workflow. Tune is built to address this, demonstrating an efficient and scalable solution for this pain point.\n",
    "\n",
    "This tutorial will walk you through the following process:\n",
    "\n",
    "1. Creating and training a model on a toy dataset (MNIST)\n",
    "2. Integrating Tune into your workflow by creating a trainable and running an experiment\n",
    "3. Trying out advanced features - plugging in an efficient scheduler\n",
    "4. (Optional) Try out a search algorithm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1: Creating and training an un-Tune-d model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import HTML\n",
    "\n",
    "\n",
    "from model import load_data, make_model, evaluate\n",
    "from helper import prepare_data, limit_threads\n",
    "\n",
    "import tensorflow as tf\n",
    "tf.logging.set_verbosity(tf.logging.ERROR)\n",
    "limit_threads(4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create and train a model to classify [MNIST](https://www.wikiwand.com/en/MNIST_database) digits **without tuning**, as a baseline first model. \n",
    "\n",
    "We will be creating a Convolutional Neural Network model (using [Keras](https://keras.io/)) to classify the digits. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_mnist():\n",
    "    x_train, y_train, x_test, y_test = load_data()\n",
    "    model = make_model(lr=0.1, layer_size=4)\n",
    "    model.fit(x_train, y_train, verbose=1, batch_size=256)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create our model. (This should take ~10 seconds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "model = train_mnist()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets evaluate the un-Tune-d model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "evaluate(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now quickly try out this model to see if it works as expected.\n",
    "\n",
    "We'll load the model with our trained weights. \n",
    "\n",
    "Try to write a digit into the box below - *do not rerun the cell after you write the digit*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "data = None\n",
    "HTML(open(\"input.html\").read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(tip: don't expect it to work)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "prepared_data = prepare_data(data)\n",
    "print(\"This model predicted your input as\", model.predict(prepared_data).argmax())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2: Setting up Tune"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import ray\n",
    "from ray import tune\n",
    "\n",
    "from helper import test_reporter, get_best_result, get_best_model\n",
    "from model import load_data, make_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One thing we might want to do now is find better hyperparameters so that our model trains more quickly and possibly to a higher accuracy. Let's make some minor modifications to utilize Tune. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1: Defining a Trainable to run\n",
    "\n",
    "Tune will automate and distribute your hyperparameter search with multiple **trials**. Each trial runs a user-defined Python function called a **trainable**. \n",
    "\n",
    "We define a new training function ``train_mnist_tune`` as our Trainable. A trainable function must:\n",
    "\n",
    " 1. pass in a `tune_reporter` object like the following: ``def train_mnist_tune(config, tune_reporter)``\n",
    " 2. train the model\n",
    " 3. report some metric(s) to Tune. \n",
    "\n",
    "Step (3) allows Tune to keep track of performance as the model is training. \n",
    "\n",
    "Tune works seamlessly with many frameworks, and in this example, we'll use a custom Keras Callback ``TuneCallback`` object to report metrics.\n",
    "\n",
    "\n",
    "The custom callback does the following on each batch:\n",
    "\n",
    "```python\n",
    "def on_batch_end(self, batch, logs={}):\n",
    "    \"\"\"Reports the last result and checkpoints if necessary.\"\"\"\n",
    "    ...\n",
    "    self.tune_reporter(\n",
    "        mean_accuracy=np.mean(self.last_10_results), \n",
    "        checkpoint=\"weights_tune.h5\")\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import keras\n",
    "import os\n",
    "\n",
    "from helper import TuneCallback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# FIXME: add an argument here so that it takes in reporter\n",
    "def train_mnist_tune(config, tune_reporter):  \n",
    "#########################################################\n",
    "    x_train, y_train, _, _ = load_data()\n",
    "    model = make_model(lr=config.get(\"lr\", 0.1), layer_size=config.get(\"layer_size\", 4))\n",
    "    callbacks = [TuneCallback(tune_reporter)]\n",
    "    for i in range(10):\n",
    "        model.fit(x_train, y_train, verbose=0, batch_size=256, callbacks=callbacks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This may take 30 seconds or so to run if incorrectly written\n",
    "assert test_reporter(train_mnist_tune)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Note: Call ``help(tune.Trainable)`` if you are interested in learning more about what qualifies as trainable in Tune.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(tune.Trainable)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Configure the search and run Tune\n",
    "\n",
    "Now that we have a working trainable, we want to use Tune to train it. We will use some basic Tune features for training: specifying a stopping criteria and a search space. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are two steps in this section. **TODO**: Configure ``tune.run`` with a stopping criteria using ``stop`` and a search space using ``config``.\n",
    "\n",
    "\n",
    "1) Set the stopping criteria to when ``mean_accuracy`` passes `0.8` and when ``training_iteration`` passes 50.\n",
    "\n",
    "For example, to specify that trials will be stopped whenever they report a `mean_accuracy` that is `>= 0.8`, do:\n",
    "\n",
    "```python\n",
    "stop={\"mean_accuracy\": 0.8}\n",
    "```\n",
    "\n",
    "\n",
    "2) We also want to designate a search space. We'll search over *learning rate*, which sets the step size of our model update, and *layer_size*, which corresponds\n",
    "to the size of the last layer of our convolutional neural network.\n",
    "\n",
    "For `layer_size`, Tune supports sampling parameters from user-specified lambda functions, which can be used independently or in combination with grid search. For learning rate, you can use `tune.grid_search` to specify an axis of a grid search. For example:\n",
    "\n",
    "```python\n",
    "space = {\n",
    "    \"lr\": tune.grid_search([0.001, 0.05, 0.1]),\n",
    "    \"layer_size\": tune.sample_from(lambda spec: int(np.random.uniform(16, 256))),\n",
    "}\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When you're ready, run the experiment! (this should take ~1 minute)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Feel free to ignore this; this is a workaround to help Ray deal with interrupts.\n",
    "ray.shutdown(); ray.init(ignore_reinit_error=True)\n",
    "#################################################\n",
    "\n",
    "\n",
    "# !! FIXME: Fix me with the instructions above\n",
    "space = {\n",
    "    \"lr\": tune.grid_search([0.001, 0.05, 0.1]),\n",
    "    \"layer_size\": tune.sample_from(lambda spec: int(np.random.uniform(16, 256))),\n",
    "}\n",
    "assert \"layer_size\" in space\n",
    "assert \"lr\" in space\n",
    "\n",
    "stop = {\n",
    "    \"mean_accuracy\": 0.8,\n",
    "    \"training_iteration\": 50\n",
    "}\n",
    "assert \"training_iteration\" in stop\n",
    "assert \"mean_accuracy\" in stop\n",
    "###############################################\n",
    "\n",
    "\n",
    "trials = tune.run(\n",
    "    train_mnist_tune,\n",
    "    stop=stop, \n",
    "    config=space,\n",
    "    resources_per_trial={\"cpu\": 2},\n",
    "    verbose=1\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can expect the result below to be about `0.6`, although your mileage may vary (and it's OK)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"The best result is\", get_best_result(trials, metric=\"mean_accuracy\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the accuracy is still low, similar to the accuracy of our first un-Tune-d model! In the next section, we will scale up the search and accelerate the training using a state of the art algorithm."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Note: Call ``help(tune.run)`` if you are interested in learning more about executing experiments.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 3: Scale up the search with more samples, hyperparameters, and a custom scheduler "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Hyperband` is state of the art algorithm that is used to early terminate un-promising trials. It has been shown to beat many state of the art hyperparameter optimization techniques and algorithms across a wide variety of datasets. **Tune allows users to take advantage this powerful algorithm in a couple lines of code.**\n",
    "\n",
    "See https://blog.ml.cmu.edu/2018/12/12/massively-parallel-hyperparameter-optimization/ for more information."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1) Create an Asynchronous HyperBand Scheduler (ASHA).\n",
    "```python\n",
    "from ray.tune.schedulers import AsyncHyperBandScheduler\n",
    "\n",
    "custom_scheduler = AsyncHyperBandScheduler(\n",
    "    reward_attr='mean_accuracy',\n",
    "    grace_period=5,\n",
    ")\n",
    "```\n",
    "\n",
    "*Note: Read the documentation on this step at https://ray.readthedocs.io/en/latest/tune-schedulers.html#asynchronous-hyperband or call ``help(tune.schedulers.AsyncHyperBandScheduler)`` to learn more about the Asynchronous Hyperband Scheduler*\n",
    "\n",
    "2) With this, we can afford to **increase the search space by 5x**. To do this, set the parameter `num_samples`. For example,\n",
    "\n",
    "```python\n",
    "tune.run(... num_samples=5)\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When you're ready, run the experiment! (this should take ~1 minute)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Feel free to ignore this; this is a workaround to help Ray deal with interrupts.\n",
    "ray.shutdown(); ray.init(ignore_reinit_error=True)\n",
    "#################################################\n",
    "\n",
    "# FIXME: Fix me with the instructions above\n",
    "from ray.tune.schedulers import AsyncHyperBandScheduler\n",
    "custom_scheduler = AsyncHyperBandScheduler(\n",
    "    reward_attr='mean_accuracy',\n",
    "    grace_period=5,\n",
    ")\n",
    "###############################################\n",
    "\n",
    "\n",
    "better_trials = tune.run(\n",
    "    train_mnist_tune,\n",
    "    # FIXME: Fix me with the instructions above\n",
    "    num_samples=5,  \n",
    "    ###############################################\n",
    "    scheduler=custom_scheduler,\n",
    "    stop={\"mean_accuracy\": 0.95, \"training_iteration\": 50},  # The mean accuracy is increased now.\n",
    "    config=space,  # We use the same search space as before.\n",
    "    resources_per_trial={\"cpu\": 2},\n",
    "    verbose=1,\n",
    "    reuse_actors=True,  # This is an internal optimization to cache actors.\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can expect the result to be about `0.95`, although your mileage may vary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(\"The best result is\", get_best_result(better_trials, metric=\"mean_accuracy\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try out a newly tuned model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tuned_model = get_best_model(make_model, better_trials, \"mean_accuracy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_data = None\n",
    "HTML(open(\"input_final.html\").read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prepared_data = prepare_data(final_data)\n",
    "print(\"This model predicted your input as\", tuned_model.predict(prepared_data).argmax())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🎉 Congratulations, you're now a Tune expert! 🎉"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Please: fill out this form to provide feedback on this tutorial!\n",
    "\n",
    "https://goo.gl/forms/NVTFjUKFz4TH8kgK2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (Optional) Try using a search algorithm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tune is an execution layer, so we can combine powerful optimizers such as HyperOpt (https://github.com/hyperopt/hyperopt) with state-of-the-art algorithms such as HyperBand without modifying any model training code.\n",
    "\n",
    "The documentation to doing this is here: https://ray.readthedocs.io/en/latest/tune-searchalg.html#hyperopt-search-tree-structured-parzen-estimators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from hyperopt import hp\n",
    "from ray.tune.suggest.hyperopt import HyperOptSearch\n",
    "\n",
    "space = {\n",
    "    \"lr\": hp.uniform(\"lr\", 0.001, 0.1),\n",
    "    \"momentum\": hp.uniform(\"momentum\", 0.1, 0.9),\n",
    "    \"hidden\": hp.choice(\"hidden\", np.arange(16, 256, dtype=int)),\n",
    "}\n",
    "\n",
    "hyperband = AsyncHyperBandScheduler(time_attr='timesteps_total', reward_attr='mean_accuracy')\n",
    "\n",
    "hyperopt_search = HyperOptSearch(space, reward_attr=\"mean_accuracy\")\n",
    "\n",
    "good_results = tune.run(\n",
    "    train_mnist_tune,\n",
    "    num_samples=5,\n",
    "    search_alg=hyperopt_search,\n",
    "    scheduler=hyperband,\n",
    "    stop={\"mean_accuracy\": 0.95},\n",
    "    config=space,\n",
    "    verbose=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(\"The best result is\", get_best_result(good_results, metric=\"mean_accuracy\"))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
